import math
import turtle
import random
from functools import partial

g_screen = None
g_snake = None  # snake's head
g_monster = None
g_monster2 = None
g_monster3 = None
g_monster4 = None
g_snake_sz = 5
# size of the snake's tail
g_intro = None
g_key_before = None
g_key_pressed = None
key_count = 0
g_status = None
running = True
snake_running = False
num_arr = []
pos_arr = []
remove_arr = set()
contact = 0
time = 0
COLOR_BODY = ("blue", "black")
COLOR_HEAD = "red"
COLOR_MONSTER = "purple"
FONT_INTRO = ("Arial", 16, "normal")
FONT_STATUS = ("Arial", 20, "normal")
FONT_GAME_OVER = ("Arial", 27, "normal")
TIMER_SNAKE = 300  # refresh rate for snake
SZ_SQUARE = 20  # square size in pixels

DIM_PLAY_AREA = 500
DIM_STAT_AREA = 40  # !!! multiple of 40
DIM_MARGIN = 30

KEY_UP, KEY_DOWN, KEY_LEFT, KEY_RIGHT, KEY_SPACE = \
    "Up", "Down", "Left", "Right", "space"

HEADING_BY_KEY = {KEY_UP: 90, KEY_DOWN: 270, KEY_LEFT: 180, KEY_RIGHT: 0}


def create_turtle(x, y, color="red", border="black"):
    """
    Creates a new turtle object with the specified position, color, and border.
    
    Args:
        x (int): The x-coordinate of the turtle's initial position.
        y (int): The y-coordinate of the turtle's initial position.
        color (str, optional): The color of the turtle's body. Defaults to "red".
        border (str, optional): The color of the turtle's border. Defaults to "black".
    
    Returns:
        turtle.Turtle: The newly created turtle object.
    """
    t = turtle.Turtle("square")
    t.color(border, color)
    t.up()
    t.goto(x, y)
    return t


def configure_play_area():
    """
    Configures the play area for the snake game, 
    including the motion border, status border, 
    introduction text, and status text.
    The motion border and status border are based on square shape
    resized according to the specified dimensions.

    Returns:
        tuples: A tuple containing the introduction text turtle and 
        the status text turtle.
    """

    # motion border
    m = create_turtle(0, 0, "", "black")
    sz = DIM_PLAY_AREA // SZ_SQUARE
    m.shapesize(sz, sz, 3)
    m.goto(0, -DIM_STAT_AREA // 2)  # shift down half the status

    # status border
    s = create_turtle(0, 0, "", "black")
    sz_w, sz_h = DIM_STAT_AREA // SZ_SQUARE, DIM_PLAY_AREA // SZ_SQUARE
    s.shapesize(sz_w, sz_h, 3)
    s.goto(0, DIM_PLAY_AREA // 2)  # shift up half the motion

    # turtle to write introduction
    intro = create_turtle(-200, 0)
    intro.hideturtle()
    # intro.write("Click anywhere to start the game .....\n\n" + \
    #             "Use arrow keys to control the snake.\n\n" + \
    #             "Press a number from 1 to 5 to extend the snake's tail.\n\n", \
    #             font=FONT_INTRO)
    intro.write("Snake by Kinley\n\n" + \
                "Click anywhere to start the game.hava fun!!\n\n",
                font=FONT_INTRO)

    # turtle to write status
    status = create_turtle(0, 0, "", "black")
    status.hideturtle()
    status.goto(-210, s.ycor() - 10)

    return intro, status


def check_collision():
    """
    Checks if the snake head has collided with the boundaries of the play area.
    If a collision is detected, the snake's movement in the direction of the collision
    is temporarily paused, allowing movement in other directions.
    """
    global snake_running
    x, y = g_snake.position()
    heading = g_snake.heading()

    next_x, next_y = x, y
    # print(heading)
    # if snake_running:
    #     if heading == 0:
    #         next_x += SZ_SQUARE
    #     elif heading == 90:
    #         next_y += SZ_SQUARE
    #     elif heading == 180:
    #         next_x -= SZ_SQUARE
    #     elif heading == 270:
    #         next_y -= SZ_SQUARE
    # print(next_x, next_y, DIM_PLAY_AREA)
    next_x = round(next_x)
    next_y = round(next_y)
    # print(next_x, next_y)
    x1 = (-DIM_PLAY_AREA / 2 + 10)
    x2 = (DIM_PLAY_AREA / 2 - 10)
    y1 = (-DIM_PLAY_AREA / 2 - 10)
    y2 = (DIM_PLAY_AREA / 2 - 30)
    if next_x >= x2 or next_x <= x1 or next_y >= y2 or next_y <= y1:
        print("check", next_x, next_y)
        # if heading == 0:
        #     g_snake.setx(x - SZ_SQUARE)
        # elif heading == 90:
        #     g_snake.sety(y - SZ_SQUARE)
        # elif heading == 180:
        #     g_snake.setx(x + SZ_SQUARE)
        # elif heading == 270:
        #     g_snake.sety(y + SZ_SQUARE)
        if next_y == y1 and g_key_pressed != KEY_DOWN and x1 < next_x < x2:
            snake_running = True
            return
        if next_y == y2 and g_key_pressed != KEY_UP and x1 < next_x < x2:
            snake_running = True
            return
        if next_x == x1 and g_key_pressed != KEY_LEFT and y1 < next_y < y2:
            snake_running = True
            return
        if next_x == x2 and g_key_pressed != KEY_RIGHT and y1 < next_y < y2:
            snake_running = True
            return
        if next_x == x1 and next_y == y2 and (g_key_pressed == KEY_DOWN or g_key_pressed == KEY_RIGHT):
            snake_running = True
            return
        if next_x == x1 and next_y == y1 and (g_key_pressed == KEY_UP or g_key_pressed == KEY_RIGHT):
            snake_running = True
            return
        if next_x == x2 and next_y == y1 and (g_key_pressed == KEY_UP or g_key_pressed == KEY_LEFT):
            snake_running = True
            return
        if next_x == x2 and next_y == y2 and (g_key_pressed == KEY_DOWN or g_key_pressed == KEY_LEFT):
            snake_running = True
            return 
        # if next_y < (-DIM_PLAY_AREA / 2 - 10):
        #     snake_running = (g_key_pressed != KEY_DOWN)
        #     print(1, next_x, next_y, snake_running)
        # elif next_y > (DIM_PLAY_AREA / 2 - 30):
        #     snake_running = (g_key_pressed != KEY_UP)
        #     print(2, next_x, next_y, snake_running)
        # elif next_x < (-DIM_PLAY_AREA / 2 + 10):
        #     snake_running = (g_key_pressed != KEY_LEFT)
        #     print(3, next_x, next_y, snake_running)
        # elif next_x > (DIM_PLAY_AREA / 2 - 10):
        #     snake_running = (g_key_pressed != KEY_RIGHT)
        #     print(4, next_x, next_y, snake_running)
        # else:
        #     snake_running = False
        #     print(5, next_x, next_y, snake_running)
        # print("check_collision", snake_running)
        snake_running = False


def configure_screen():
    """
    Configures the Turtle screen for the snake game,
    the screen width and height are calculated based 
    on the play area, status bar and margin.
    
    Returns:
        turtle.Screen: The configured Turtle screen.
    """
    s = turtle.Screen()
    s.tracer(0)  # disable auto screen refresh, 0=disable, 1=enable
    s.title("Snake by Kinley")
    w = DIM_PLAY_AREA + DIM_MARGIN * 2
    h = DIM_PLAY_AREA + DIM_MARGIN * 2 + DIM_STAT_AREA
    s.setup(w, h)
    s.mode("standard")
    return s


def update_status():
    """
    Updates the status display on the screen with 
    the current key press and snake tail length.
    """
    g_status.clear()
    # Tail length-{g_snake_sz}
    motion = g_key_pressed
    if motion is None:
        motion = "Paused"
    status = f'Contact:{contact}  Time:{time}   Motion: {motion}'
    g_status.write(status, font=FONT_STATUS)
    g_screen.update()


def on_arrow_key_pressed(key):
    """
    Handles the user's arrow key press event and 
    updates the global `g_key_pressed` variable with the pressed key. 
    It then calls the `update_status()` function to update 
    the status display on the screen.
    
    Args:
        key (str): The key that was pressed, one of 'Up', 'Down', 'Left', or 'Right'.
    """
    global g_key_pressed, g_key_before, key_count, snake_running
    key_count += 1
    if key == KEY_SPACE and snake_running:
        g_key_before = g_key_pressed
    g_key_pressed = key
    if key == KEY_SPACE:
        g_key_pressed = None
        snake_running = not snake_running
        if snake_running:
            g_key_pressed = g_key_before
    
    next_x, next_y = g_snake.position()
    next_x = round(next_x)
    next_y = round(next_y)
    # print("pressed", next_x, next_y)
    if (g_key_pressed == KEY_LEFT or g_key_pressed == KEY_RIGHT) and (-DIM_PLAY_AREA / 2 + 10) < next_x < (DIM_PLAY_AREA / 2 - 10):
        snake_running = True
    if (g_key_pressed == KEY_UP or g_key_pressed == KEY_DOWN) and (-DIM_PLAY_AREA / 2 - 10) < next_y < (DIM_PLAY_AREA / 2 - 30):
        snake_running = True
    update_status()


def on_timer_snake():
    """
    Advances the snake's movement on a timer. This function is called repeatedly 
    by the Turtle screen's `ontimer` method to update the snake's position.
    
    If no key has been pressed, the function simply schedules itself to be called 
    again after the `TIMER_SNAKE` interval. Otherwise, it performs the following steps:
    
    1. Clones the snake's head as a new body segment by setting the color to `COLOR_BODY` 
        and stamping it on the screen.
    2. Sets the snake's color back to `COLOR_HEAD`.
    3. Advances the snake's position by setting the heading based 
        on the last key pressed (`g_key_pressed`) and 
        moving the snake forward by `SZ_SQUARE` units.
    4. If the number of stamped segments exceeds the current snake size (`g_snake_sz`), 
        removes the last segment by clearing the oldest stamp.
    5. Updates the Turtle screen to reflect the changes.
    6. Schedules the function to be called again after the `TIMER_SNAKE` interval.
    """
    # print(on_timer_snake)
    global running
    if not running:
        return
    if g_key_pressed is None:
        g_screen.ontimer(on_timer_snake, TIMER_SNAKE)
        return

    check_collision()
    if not snake_running:
        g_screen.ontimer(on_timer_snake, TIMER_SNAKE)
        return
    
    for n in range(1, 10):
        if (n - 1) in remove_arr:
            continue
        num = num_arr[n - 1]
        x1 = g_snake.xcor()
        y1 = g_snake.ycor()
        x2 = num.xcor()
        y2 = num.ycor()
        d = math.sqrt((x1 - x2) ** 2 + (y1 - y2) ** 2)
        #print(d)
        if d <= 20:
            remove_arr.add(n - 1)
            consume_food(n)
            num.clear()
            if len(remove_arr) == 9:
                running = False
                intro = create_turtle(-100, 0)
                intro.color("red")
                intro.hideturtle()
                intro.write("Winner!!!", font=FONT_GAME_OVER)
    

    # Clone the head as body
    g_snake.color(*COLOR_BODY)
    g_snake.stamp()
    pos_arr.append((g_snake.xcor(), g_snake.ycor()))
    g_snake.color(COLOR_HEAD)

    # Advance snake
    g_snake.setheading(HEADING_BY_KEY[g_key_pressed])
    g_snake.forward(SZ_SQUARE)

    # Shifting or extending the tail.
    # Remove the last square on Shifting.
    if len(g_snake.stampItems) > g_snake_sz:
        g_snake.clearstamps(1)
        pos_arr.pop(0)


    g_screen.update()

    g_screen.ontimer(on_timer_snake, TIMER_SNAKE)


def on_timer_food():
    global running
    if not running:
        return
    for n in range(1, 10):
        num = num_arr[n - 1]
        num.clear()
        if (n - 1) in remove_arr:
            continue
        head_arr = []
        x, y = num.position()
        if (y - SZ_SQUARE * 2) > (-DIM_PLAY_AREA / 2):
            head_arr.append(KEY_DOWN)
        elif (y + SZ_SQUARE * 2) < (DIM_PLAY_AREA / 2 - 30):
            head_arr.append(KEY_UP)
        elif (x - SZ_SQUARE * 2) > (-DIM_PLAY_AREA / 2 + 10):
            head_arr.append(KEY_LEFT)
        elif (x + SZ_SQUARE * 2) < (DIM_PLAY_AREA / 2 - 10):
            head_arr.append(KEY_RIGHT)
        heading = head_arr[random.randint(0, len(head_arr) - 1)]
        num.setheading(HEADING_BY_KEY[heading])
        num.forward(SZ_SQUARE * 2)
        # print(x, num.ycor())
        
        num.write(n, font=FONT_INTRO)
    g_screen.ontimer(on_timer_food, 6000)


def on_timer_monster(monster):
    """
    Advances the monster's movement on a timer. 
    This function is called repeatedly 
    by the Turtle screen's `ontimer` method to update the monster's position.
    
    The function performs the following steps:
    
    1. Calculates the heading for the monster to move towards the snake, 
        snapping to the nearest 45-degree angle.
    2. Sets the monster's heading to the calculated value.
    3. Moves the monster forward by `SZ_SQUARE` units.
    4. Updates the Turtle screen to reflect the changes.
    5. Schedules the function to be called again after a random delay 
        between `TIMER_SNAKE-50` and `TIMER_SNAKE+200` milliseconds.
    """
    global running, g_intro, contact
    x1 = g_snake.xcor()
    y1 = g_snake.ycor()
    x2 = monster.xcor()
    y2 = monster.ycor()
    d = math.sqrt((x1 - x2) ** 2 + (y1 - y2) ** 2)
    if d <= 20:
        running = False
        intro = create_turtle(-100, 0)
        intro.color("red")
        intro.hideturtle()
        intro.write("Game Over!!!", font=FONT_GAME_OVER)
    for stamp_pos in pos_arr:
        dd = math.sqrt((stamp_pos[0] - x2) ** 2 + (stamp_pos[1] - y2) ** 2)
        if dd <= 20:
            contact += 1
            update_status()
    if not running:
        return
    angle = monster.towards(g_snake)
    qtr = angle // 45  # (0,1,2,3,4,5,6,7)
    heading = qtr * 45 if qtr % 2 == 0 else (qtr + 1) * 45

    monster.setheading(heading)
    monster.forward(SZ_SQUARE)

    g_screen.update()
    delay = random.randint(TIMER_SNAKE + 100, TIMER_SNAKE + 1000)

    # func = function.partial(on_timer_monster, monster)
    def func():
        on_timer_monster(monster)

    g_screen.ontimer(func, delay)



def on_timer_count():
    if not running:
        return
    global time
    time += 1
    update_status()
    g_screen.ontimer(on_timer_count, 1000)


def consume_food(num):
    """
    Simulate the action when the snake consumes the given food item num.
    Increase the length of the snake's tail by the given num value.
    
    Args:
        num (int): The value of food item being consumed by the snake
    
    Modifies:
        g_snake_sz (int): Increments the size of the snake's body by the specified number.
        update_status(): Updates the game status display to reflect the new snake size.
    """

    global g_snake_sz
    if g_snake_sz < 100:
        g_snake_sz += num
        update_status()


def cb_start_game(x, y):
    """
    Starts the game by setting up the initial game state and event handlers.
    
    This function is called when the user clicks on the screen to start the game.
    It performs the following steps:
    
    1. Clears the on-screen click handler to prevent further clicks from starting the game.
    2. Clears the introductory message.
    3. Registers key event handlers for the arrow keys to handle player movement.
    4. Starts the timer-based updates for the snake and monster movements.
    5. Registers key event handlers for consuming food items 1 through 5.
    """
    global snake_running
    g_screen.onscreenclick(None)
    g_intro.clear()

    for key in (KEY_UP, KEY_DOWN, KEY_RIGHT, KEY_LEFT, KEY_SPACE):
        g_screen.onkey(partial(on_arrow_key_pressed, key), key)

    for n in range(1, 6):
        num_arr[n - 1].write(n, font=FONT_INTRO)
        g_screen.onkey(partial(consume_food, n), n)

    snake_running = True
    on_timer_count()
    on_timer_snake()
    on_timer_food()
    on_timer_monster(g_monster)
    on_timer_monster(g_monster2)
    on_timer_monster(g_monster3)
    on_timer_monster(g_monster4)


if __name__ == "__main__":
    """
    A small demo illustrating the main processing logic 
    for a Snake game using Python Turtle graphics.
    """
    g_screen = configure_screen()
    g_intro, g_status = configure_play_area()

    update_status()

    g_monster = create_turtle(-110, -110, COLOR_MONSTER, "black")
    g_monster2 = create_turtle(110, 110, COLOR_MONSTER, "black")
    g_monster3 = create_turtle(-110, 110, COLOR_MONSTER, "black")
    g_monster4 = create_turtle(110, -110, COLOR_MONSTER, "black")

    for n in range(1, 10):
        intro = create_turtle(random.randint(-250 + 20, 250 - 20), random.randint(-250 + 50, 250 - 50))
        intro.hideturtle()
        num_arr.append(intro)

    g_snake = create_turtle(0, 0, COLOR_HEAD, "black")

    g_screen.onscreenclick(cb_start_game)  # set up a mouse-click call back

    g_screen.update()
    g_screen.listen()
    g_screen.mainloop()
